import numpy as np
from enum import Enum
from typing import Dict

class SubChannelType(Enum):
    """子通道类型枚举"""
    CORNER = "corner"    # 角通道
    EDGE = "edge"       # 边通道
    CENTER = "center"   # 中心通道

class SubChannel:
    """子通道类，定义单个子通道的几何和物理特性"""
    def __init__(self, id: int, type: SubChannelType, geometry: Dict):
        self.id = id
        self.type = type
        
        # 几何参数
        self.area = geometry['area']                          # 流通面积
        self.wetted_perimeter = geometry['wetted_perimeter']  # 湿周长
        self.heated_perimeter = geometry['heated_perimeter']  # 加热周长
        self.hydraulic_diameter = geometry['hydraulic_diameter']  # 水力直径
        self.gap_width = geometry['gap_width']                # 间隙宽度
        self.gap_length = geometry['gap_length']              # 间隙长度
        
        # 相邻子通道连接信息
        self.connections = []  # 存储相邻子通道的连接信息
        
        # 场量初始化
        self.alpha = None              # 空泡份额
        self.pressure = None           # 压力场
        self.velocity = {              # 速度场
            'liquid': {
                'axial': None,
                'lateral': None
            },
            'vapor': {
                'axial': None,
                'lateral': None
            }
        }
        
        # 保存上一时间步的值
        self.previous_alpha = None
        self.previous_pressure = None
        self.previous_velocity = {     # 上一时间步速度场
            'liquid': {
                'axial': None,
                'lateral': None
            },
            'vapor': {
                'axial': None,
                'lateral': None
            }
        }
        
        # 温度场
        self.temperature = {
            'liquid': None,
            'vapor': None
        }
        
        # 密度场
        self.density = {
            'liquid': None,
            'vapor': None
        }
        
        # TVD格式相关
        self.previous_drift_flux = None  # 上一时间步的漂移通量
        self.previous_mass_flux = {      # 上一时间步的质量通量
            'vapor': None,
            'liquid': None
        }
    
    def add_connection(self, neighbor_id: int, connection_type: str, gap_width: float):
        """添加与相邻子通道的连接信息
        
        Args:
            neighbor_id: 相邻子通道的ID
            connection_type: 连接类型 ('axial' 或 'lateral')
            gap_width: 间隙宽度
        """
        self.connections.append({
            'neighbor_id': neighbor_id,
            'type': connection_type,
            'gap_width': gap_width
        })
    
    def get_connection_parameters(self, adjacent_channel) -> Dict:
        """获取与相邻子通道的连接参数
        
        Args:
            adjacent_channel: 相邻子通道对象
            
        Returns:
            Dict: 包含连接参数的字典
        """
        # 查找连接信息
        connection = next(
            (conn for conn in self.connections 
             if conn['neighbor_id'] == adjacent_channel.id),
            None
        )
        
        if connection is None:
            raise ValueError(f"未找到与通道 {adjacent_channel.id} 的连接信息")
        
        # 计算质心距离（简化模型，使用固定值）
        if connection['type'] == 'lateral':
            centroid_distance = self.gap_width
        else:  # axial
            centroid_distance = self.gap_length
        
        return {
            'connection_type': connection['type'],
            'gap_width': connection['gap_width'],
            'gap_length': self.gap_length,
            'centroid_distance': centroid_distance
        }
    
    def get_geometry_parameters(self) -> Dict:
        """获取子通道的几何参数
        
        Returns:
            Dict: 包含所有几何参数的字典
        """
        return {
            'area': self.area,
            'wetted_perimeter': self.wetted_perimeter,
            'heated_perimeter': self.heated_perimeter,
            'hydraulic_diameter': self.hydraulic_diameter,
            'gap_width': self.gap_width,
            'gap_length': self.gap_length
        }
    
    def save_current_state(self):
        """保存当前状态用于下一时间步计算"""
        self.previous_alpha = np.copy(self.alpha)
        self.previous_pressure = np.copy(self.pressure)
        self.previous_velocity = {
            'vapor': {
                'axial': np.copy(self.velocity['vapor']['axial']),
                'lateral': np.copy(self.velocity['vapor']['lateral'])
            },
            'liquid': {
                'axial': np.copy(self.velocity['liquid']['axial']),
                'lateral': np.copy(self.velocity['liquid']['lateral'])
            }
        }
        
        # 初始化或更新TVD相关数组
        if self.previous_drift_flux is None:
            self.previous_drift_flux = np.zeros_like(self.alpha)
        if self.previous_mass_flux['vapor'] is None:
            self.previous_mass_flux['vapor'] = np.zeros_like(self.alpha)
            self.previous_mass_flux['liquid'] = np.zeros_like(self.alpha)

class Grid:
    """网格类，管理所有子通道和其连接关系"""
    def __init__(self, n_axial=20, channel_length=3.66, cos_theta=1.0):
        """网格类构造函数

        Args:
            n_axial (int, optional): 轴向节点数. Defaults to 20.
            channel_length (float, optional): 通道长度. Defaults to 3.66.
            cos_theta (float, optional): 通道与重力方向夹角的余弦值. Defaults to 1.0.
        """
        self.n_axial = n_axial  # 轴向节点数
        self.n_channels = 36     # 5x5棒束的子通道总数
        self.subchannels = []    # 存储所有子通道
        self.channel_length = channel_length  # 通道长度
        self.dz = channel_length / (n_axial - 1)  # 轴向网格间距
        self.dx = 0.013  # 横向网格间距（根据燃料棒间距设置）
        self.cos_theta = cos_theta # NOTE: 通道与重力方向夹角的余弦值
        
        # 时间步相关
        self.current_time = 0.0  # 当前时间
        self.dt = 0.0  # 当前时间步长
        self.time_step_number = 0  # 当前时间步编号
        
        # 网格索引映射
        self.axial_indices = np.arange(n_axial)
        self.channel_indices = np.arange(self.n_channels)

        # 壁面通道ID列表
        self.wall_channel_ids = []
        
        # 子通道类型映射
        self.channel_types = self._initialize_channel_types()
        
        # 创建所有子通道
        self._create_all_subchannels()
        
        # 建立子通道之间的连接关系
        self.setup_connections()
        
        # 初始化场变量
        self.initialize_field_variables()
        
    def _initialize_channel_types(self):
        """初始化5x5棒束的子通道类型映射"""
        # 5x5棒束有36个子通道：
        # - 4个角通道
        # - 16个边通道
        # - 16个中心通道
        channel_types = {}

        # 角通道 (4个)
        corner_channels = [0, 5, 30, 35]
        for id in corner_channels:
            channel_types[id] = SubChannelType.CORNER

        # 边通道 (16个)
        edge_channels = [
            1, 2, 3, 4,         # 上边
            31, 32, 33, 34,    # 下边
            6, 12, 18, 24,    # 左边
            11, 17, 23, 29     # 右边
        ]
        for id in edge_channels:
            channel_types[id] = SubChannelType.EDGE
            
        # 中心通道 (16个)
        center_channels = [i for i in range(self.n_channels) 
                         if i not in corner_channels + edge_channels]
        for id in center_channels:
            channel_types[id] = SubChannelType.CENTER

        # NOTE: 任务 2.1 - 记录所有壁面通道（角通道和边通道）
        self.wall_channel_ids = corner_channels + edge_channels
            
        return channel_types
        
    def create_subchannel(self, id):
        """创建新的子通道，根据ID确定类型和几何参数"""
        channel_type = self.channel_types[id % self.n_channels]
        
        # 根据COBRA-EN的经验计算几何参数
        rod_diameter = 0.01    # 燃料棒直径
        rod_pitch = 0.013     # 燃料棒间距
        
        if channel_type == SubChannelType.CORNER:
            # 角通道
            area = 0.25 * np.pi * (rod_pitch - rod_diameter)**2
            wetted_perimeter = 0.5 * np.pi * rod_diameter + 2 * (rod_pitch - rod_diameter)
            heated_perimeter = 0.5 * np.pi * rod_diameter
            gap_width = rod_pitch - rod_diameter
            gap_length = rod_pitch
            
        elif channel_type == SubChannelType.EDGE:
            # 边通道
            area = (rod_pitch - rod_diameter) * rod_pitch - 0.25 * np.pi * rod_diameter**2
            wetted_perimeter = 0.5 * np.pi * rod_diameter + 2 * rod_pitch
            heated_perimeter = 0.5 * np.pi * rod_diameter
            gap_width = rod_pitch - rod_diameter
            gap_length = rod_pitch
            
        else:  # CENTER
            # 中心通道
            area = rod_pitch**2 - np.pi * rod_diameter**2 / 4
            wetted_perimeter = np.pi * rod_diameter
            heated_perimeter = np.pi * rod_diameter
            gap_width = rod_pitch - rod_diameter
            gap_length = rod_pitch
        
        # 计算水力直径
        hydraulic_diameter = 4 * area / wetted_perimeter
        
        geometry = {
            'area': area,
            'wetted_perimeter': wetted_perimeter,
            'heated_perimeter': heated_perimeter,
            'hydraulic_diameter': hydraulic_diameter,
            'gap_width': gap_width,
            'gap_length': gap_length
        }
        
        subchannel = SubChannel(id, channel_type, geometry)
        self.subchannels.append(subchannel)
        return subchannel
    
    def setup_connections(self):
        """建立5×5棒束36子通道之间的连接关系"""
        # 定义单层内的横向连接关系模板（基于36个通道的拓扑结构）
        lateral_connections_template = {
            # 角通道连接 (4个)
            0: [(1, 'lateral'), (6, 'lateral')],
            5: [(4, 'lateral'), (11, 'lateral')],
            30: [(24, 'lateral'), (31, 'lateral')],
            35: [(29, 'lateral'), (34, 'lateral')],
            
            # 边通道连接 (8个)
            1: [(0, 'lateral'), (2, 'lateral'), (7, 'lateral')],
            2: [(1, 'lateral'), (3, 'lateral'), (8, 'lateral')],
            3: [(2, 'lateral'), (4, 'lateral'), (9, 'lateral')],
            4: [(3, 'lateral'), (5, 'lateral'), (10, 'lateral')],
            31: [(30, 'lateral'), (32, 'lateral'), (25, 'lateral')],
            32: [(31, 'lateral'), (33, 'lateral'), (26, 'lateral')],
            33: [(32, 'lateral'), (34, 'lateral'), (27, 'lateral')],
            34: [(33, 'lateral'), (35, 'lateral'), (28, 'lateral')],
            
            # 中心通道连接 (24个)
            6: [(0, 'lateral'), (7, 'lateral'), (12, 'lateral')],
            7: [(1, 'lateral'), (6, 'lateral'), (8, 'lateral'), (13, 'lateral')],
            8: [(2, 'lateral'), (7, 'lateral'), (9, 'lateral'), (14, 'lateral')],
            9: [(3, 'lateral'), (8, 'lateral'), (10, 'lateral'), (15, 'lateral')],
            10: [(4, 'lateral'), (9, 'lateral'), (11, 'lateral'), (16, 'lateral')],
            11: [(5, 'lateral'), (10, 'lateral'), (17, 'lateral')],
            
            12: [(6, 'lateral'), (13, 'lateral'), (18, 'lateral')],
            13: [(7, 'lateral'), (12, 'lateral'), (14, 'lateral'), (19, 'lateral')],
            14: [(8, 'lateral'), (13, 'lateral'), (15, 'lateral'), (20, 'lateral')],
            15: [(9, 'lateral'), (14, 'lateral'), (16, 'lateral'), (21, 'lateral')],
            16: [(10, 'lateral'), (15, 'lateral'), (17, 'lateral'), (22, 'lateral')],
            17: [(11, 'lateral'), (16, 'lateral'), (23, 'lateral')],
            
            18: [(12, 'lateral'), (19, 'lateral'), (24, 'lateral')],
            19: [(13, 'lateral'), (18, 'lateral'), (20, 'lateral'), (25, 'lateral')],
            20: [(14, 'lateral'), (19, 'lateral'), (21, 'lateral'), (26, 'lateral')],
            21: [(15, 'lateral'), (20, 'lateral'), (22, 'lateral'), (27, 'lateral')],
            22: [(16, 'lateral'), (21, 'lateral'), (23, 'lateral'), (28, 'lateral')],
            23: [(17, 'lateral'), (22, 'lateral'), (29, 'lateral')],
            
            24: [(18, 'lateral'), (25, 'lateral'), (30, 'lateral')],
            25: [(19, 'lateral'), (24, 'lateral'), (26, 'lateral'), (31, 'lateral')],
            26: [(20, 'lateral'), (25, 'lateral'), (27, 'lateral'), (32, 'lateral')],
            27: [(21, 'lateral'), (26, 'lateral'), (28, 'lateral'), (33, 'lateral')],
            28: [(22, 'lateral'), (27, 'lateral'), (29, 'lateral'), (34, 'lateral')],
            29: [(23, 'lateral'), (28, 'lateral'), (35, 'lateral')]
        }
        
        # FIXME: 修复了连接逻辑。轴向连接是隐式的（通过数组索引），只需要建立横向连接。
        lateral_connections_count = 0
        for channel_id, conn_list in lateral_connections_template.items():
            for neighbor_id, conn_type in conn_list:
                # 确保ID在有效范围内
                if 0 <= channel_id < self.n_channels and 0 <= neighbor_id < self.n_channels:
                    # FIXME: 任务 1.3 - 使用正确的几何参数 gap_width
                    # 从子通道对象中获取正确的间隙宽度，而不是使用全局的dx
                    channel = self.subchannels[channel_id]
                    neighbor = self.subchannels[neighbor_id]
                    # 通常，两个通道间的gap是共享的，这里简化为使用当前通道的gap_width
                    gap_width = channel.gap_width
                    self.subchannels[channel_id].add_connection(neighbor_id, conn_type, gap_width)
                    lateral_connections_count += 1
        
        print(f"建立了{lateral_connections_count}个横向连接")
    
    def initialize_field_variables(self):
        """初始化场变量"""
        for subchannel in self.subchannels:
            # 初始化空泡份额
            subchannel.alpha = np.zeros(self.n_axial)
            subchannel.previous_alpha = np.zeros(self.n_axial)
            
            # 初始化压力场
            subchannel.pressure = np.ones(self.n_axial) * 1.0e5  # 假设初始压力为1bar
            subchannel.previous_pressure = np.ones(self.n_axial) * 1.0e5
            
            # 初始化速度场
            for phase in ['liquid', 'vapor']:
                subchannel.velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['axial'] = np.zeros(self.n_axial)
                subchannel.previous_velocity[phase]['lateral'] = np.zeros(self.n_axial)
                subchannel.temperature[phase] = np.ones(self.n_axial) * 300.0  # 假设初始温度为300K
                subchannel.density[phase] = np.zeros(self.n_axial)
    
    def set_boundary_conditions(self, inlet_conditions, outlet_conditions):
        """设置边界条件"""
        # FIXME: 修复了边界条件设置逻辑，以适应新的网格结构。
        # 入口边界条件
        for subchannel in self.subchannels:  # 遍历所有36个子通道
            # 设置速度
            for phase in ['liquid', 'vapor']:
                # 从inlet_conditions中获取正确的速度值
                subchannel.velocity[phase]['axial'][0] = inlet_conditions['velocity'][phase]['axial']
                # 使用全局温度
                subchannel.temperature[phase][0] = inlet_conditions['temperature']
            
            # 设置压力和空泡份额
            subchannel.pressure[0] = inlet_conditions['pressure']
            subchannel.alpha[0] = inlet_conditions['alpha']
        
        # 出口边界条件
        for subchannel in self.subchannels:  # 遍历所有36个子通道
            subchannel.pressure[-1] = outlet_conditions['pressure']
            
    def advance_time(self, dt):
        """推进时间步"""
        self.current_time += dt
        self.time_step_number += 1
        self.dt = dt
        
        # 保存当前状态为上一时间步的值
        for subchannel in self.subchannels:
            subchannel.save_current_state() 
    
    def is_wall_channel(self, channel_id: int) -> bool:
        """判断一个子通道是否是壁面通道

        Args:
            channel_id (int): 子通道ID

        Returns:
            bool: 如果是壁面通道则返回True
        """
        return channel_id in self.wall_channel_ids

    def _create_all_subchannels(self):
        """创建所有子通道实例"""
        # FIXME: 修复了网格创建逻辑。之前错误地为每个轴向节点都创建了子通道对象，
        # 正确的做法是只创建一次子通道，每个子通道内部包含所有轴向节点的数据。
        self.subchannels = []
        for i in range(self.n_channels):
            self.create_subchannel(i)
    
